from gurobipy import *
import numpy as np
from traffic_engineer import *
from traffic_engineering_class import *
import copy

# weight each path greedily, so that weights of all paths between two points sum to 1
# splits traffic among paths based on the relative capacity of the paths. Essentially VLB
class ObliviousRouting(TrafficEngineer):
	def __init__(self, aurora_network, all_interblock_paths):
		TrafficEngineer.__init__(self, aurora_network, all_interblock_paths)
		return

	def get_filename_param(self):
		return "vlb"

	# returns a string description of this class
	def get_string(self):
		return "Oblivious Routing"

	def traffic_engineering_class(self):
		return TrafficEngineeringClass.WCMP

	#def _compute_mlu(self, traffic_matrices, capacity_matrix, all_paths):
	## Currently does not involve another step that reduces hop count
	def _compute_path_weights_minimize_worstcase_mlu(self, logical_topology_adj_matrix, traffic_matrices):
		nblocks = self.aurora_network.get_num_blocks()
		num_tm = len(traffic_matrices)

		max_tm = np.zeros((nblocks, nblocks))
		for tm in traffic_matrices:
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						max_tm[i][j] = max(tm[i][j], max_tm[i][j])

		model = Model("Routing minimize MLU for all traffic matrices")
		model.setParam( 'OutputFlag', False )
		mlu = model.addVar(lb=0., ub=GRB.INFINITY, obj=0, vtype=GRB.CONTINUOUS, name="mlu")
		routing_weight_vars = {}
		link_capacity_constraints = [None] * num_tm
		for k in range(num_tm):
			link_capacity_constraints[k] = [None] * nblocks
			for i in range(nblocks):
				link_capacity_constraints[k][i] = [None] * nblocks
				for j in range(nblocks):
					if i != j:
						link_capacity_constraints[k][i][j] = LinExpr()
						#routing_weight_vars[path] = model.addVar(lb=0., ub=1., obj=1., vtype=GRB.CONTINUOUS, name="w_{},{}".format(i,j))
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					weight_constraint = LinExpr()
					for path in self.all_interblock_paths[i][j]:
						routing_weight_vars[path] = model.addVar(lb=0, ub=1., obj=1., vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						weight_constraint += routing_weight_vars[path]
					model.addConstr(lhs=weight_constraint, sense=GRB.EQUAL, rhs=1.)


		# stage 1: setup all the optimization variables that are the routing weights
		# stage 1.5 : also adds the traffic flow satisfiability constraints
		for tm, tm_index in zip(traffic_matrices, range(num_tm)):
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						for path in self.all_interblock_paths[i][j]:
							src = path[0]
							for path_hop in range(1, len(path), 1):
								dst = path[path_hop]
								link_capacity_constraints[tm_index][src][dst] += (routing_weight_vars[path] * tm[i][j])
								src = dst

		## stage 2 : add the link utilization constraints to the model
		for tm_index in range(num_tm):
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						lu_constraint = link_capacity_constraints[tm_index][i][j]
						capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
						model.addConstr(lhs=lu_constraint, sense=GRB.LESS_EQUAL, rhs=mlu * capacity * logical_topology_adj_matrix[i][j])

		# stage 3: set the objective
		model.setObjective(mlu, GRB.MINIMIZE)
		try: 
			model.optimize()
			routing_weights = {}
			for path in routing_weight_vars.keys():
				routing_weights[path] = routing_weight_vars[path].x
			return routing_weights, mlu.x
		except GurobiError as e:
			print ("Error code " + str(e. errno ) + ": " + str(e))
			return None
		except AttributeError :
			print ("Encountered an attribute error in Worst case MLU")
			return None

	def __recursive_pairs(self, current_pod_id, nblocks, leftover_options, pairs_so_far, all_sets):
		if current_pod_id >= nblocks and len(all_sets) == nblocks:
			all_sets.append(pairs_so_far)
			return
		for target_pod in leftover_options:
			if target_pod != current_pod_id:
				pairs_so_far_current = copy.deepcopy(pairs_so_far)
				leftover_options_current = copy.deepcopy(leftover_options)
				pairs_so_far.append((current_pod_id, target_pod))
				leftover_options_current.remove(target_pod)
				self.__recursive_pairs(current_pod_id + 1, nblocks, leftover_options_current, pairs_so_far_current, all_sets)
		return

	def _form_permutation_traffic_matrices(self, nblocks):
		all_sets = []
		leftover_pods_set = set()
		for i in range(nblocks):
			leftover_pods_set.add(i)
		self.__recursive_pairs(0, nblocks, leftover_pods_set, [], all_sets)

		capacity_matrix = np.zeros((nblocks, nblocks))
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					capacity_matrix[i][j] = min(self.aurora_network.get_num_links(i) * self.aurora_network.get_link_capacity(i), self.aurora_network.get_num_links(j) * self.aurora_network.get_link_capacity(j))

		permutation_matrices = []
		for perm_matrix_setting in all_sets:
			tm = np.zeros((nblocks, nblocks))
			assert(len(perm_matrix_setting) == nblocks)
			for i, j in perm_matrix_setting:
				tm[i][j] = capacity_matrix
			permutation_matrices.append(tm)
		return permutation_matrices

	def compute_path_weights(self, logical_topology_adj_matrix, traffic_matrices):
		nblocks = self.aurora_network.get_num_blocks()

		permutation_matrices = self._form_permutation_traffic_matrices(nblocks)
		routing_weights, mlu = self._compute_path_weights_minimize_worstcase_mlu(logical_topology_adj_matrix, permutation_matrices)
		return routing_weights

